package com.aiwiown.snackmq.management.scaling;

import com.aiwiown.snackmq.management.BrokerMetrics;
import lombok.extern.slf4j.Slf4j;

import java.time.Instant;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 自动调节服务，根据策略和指标动态调整系统行为。
 */
@Slf4j
public class AutoScalerService implements AutoCloseable {

    private final ScalingPolicy policy;
    private final BrokerMetrics metrics;
    private final ScheduledExecutorService scheduler;
    private final AtomicBoolean throttlingEnabled = new AtomicBoolean(false);
    private volatile Instant lastActionTimestamp = Instant.MIN;

    // --- 新增字段用于计算速率 ---
    private long lastCheckedCount = 0;
    private long lastCheckedTimestamp = -1;

    /**
     * 构造函数，初始化自动调节服务。
     *
     * @param policy  调节策略对象，包含高 TPS 阈值、低 TPS 阈值、检查间隔等配置。
     * @param metrics 消息代理的指标对象，用于获取消息生产总数等数据。
     */
    public AutoScalerService(ScalingPolicy policy, BrokerMetrics metrics) {
        this.policy = policy;
        this.metrics = metrics;
        this.scheduler = Executors.newSingleThreadScheduledExecutor(
                r -> new Thread(r, "snackmq-auto-scaler")
        );
    }

    /**
     * 启动自动调节服务，按照策略中定义的检查间隔定期评估系统状态。
     * 如果策略未启用，则不会启动调度器。
     */
    public void start() {
        if (!policy.isEnabled()) {
            log.info("Auto-scaling is disabled by policy.");
            return;
        }
        log.info("Starting AutoScalerService with check interval of {} seconds.", policy.getCheckIntervalSeconds());

        // 初始化首次检查的时间戳和计数值
        this.lastCheckedTimestamp = System.nanoTime();
        this.lastCheckedCount = metrics.getTotalMessagesProduced();

        scheduler.scheduleAtFixedRate(this::evaluate,
                policy.getCheckIntervalSeconds(),
                policy.getCheckIntervalSeconds(),
                TimeUnit.SECONDS);
    }

    /**
     * 定期评估当前系统的 TPS（每秒事务数），并根据策略决定是否启用或解除限流。
     * 该方法由调度器定期调用。
     */
    private void evaluate() {
        try {
            // --- 关键修改：手动计算速率 ---
            long currentCount = metrics.getTotalMessagesProduced();
            long currentTimestamp = System.nanoTime();

            long countDelta = currentCount - lastCheckedCount;
            double timeDeltaSeconds = (currentTimestamp - lastCheckedTimestamp) / 1_000_000_000.0;

            // 更新状态以备下次计算
            this.lastCheckedCount = currentCount;
            this.lastCheckedTimestamp = currentTimestamp;

            // 避免因时间间隔过小或为零导致除零异常
            if (timeDeltaSeconds <= 0) {
                return;
            }

            double currentTps = countDelta / timeDeltaSeconds;

            log.debug("Evaluating scaling policy. Current TPS: {}", String.format("%.2f", currentTps));

            // 检查是否需要启动限流 (Scale Down)
            if (currentTps > policy.getHighTpsThreshold()) {
                if (isCooldownActive(policy.getScaleDownCooldownSeconds())) {
                    log.debug("Scale-down cooldown is active. No action will be taken.");
                    return;
                }
                log.warn("High TPS threshold breached ({} > {}). Enabling producer throttling.",
                        String.format("%.2f", currentTps), policy.getHighTpsThreshold());
                throttlingEnabled.set(true);
                lastActionTimestamp = Instant.now();
            }
            // 检查是否需要解除限流 (Scale Up)
            else if (currentTps < policy.getLowTpsThreshold()) {
                if (isCooldownActive(policy.getScaleUpCooldownSeconds())) {
                    log.debug("Scale-up cooldown is active. No action will be taken.");
                    return;
                }
                if (throttlingEnabled.get()) {
                    log.info("TPS is below low threshold ({} < {}). Disabling producer throttling.",
                            String.format("%.2f", currentTps), policy.getLowTpsThreshold());
                    throttlingEnabled.set(false);
                    lastActionTimestamp = Instant.now();
                }
            }
        } catch (Exception e) {
            log.error("Error during auto-scaling evaluation.", e);
        }
    }

    /**
     * 检查冷却时间是否仍然有效。
     *
     * @param cooldownSeconds 冷却时间的秒数。
     * @return 如果冷却时间仍然有效，则返回 true；否则返回 false。
     */
    private boolean isCooldownActive(int cooldownSeconds) {
        return Instant.now().isBefore(lastActionTimestamp.plusSeconds(cooldownSeconds));
    }

    /**
     * 查询当前是否需要限流。
     *
     * @return 如果需要限流则返回 true，否则返回 false。
     */
    public boolean isThrottling() {
        return throttlingEnabled.get();
    }

    /**
     * 关闭自动调节服务，释放相关资源。
     */
    @Override
    public void close() {
        log.info("Shutting down AutoScalerService.");
        scheduler.shutdownNow();
    }
}
